import torch
import cv2 as cv
from model.zhnnet import ZhnNet, image_generate_conf, image_generate_loc

print('Start detect via camera.')
classify = False
model = ZhnNet(classify=classify)
model.load_state_dict(torch.load('zhnnet.pth',map_location=torch.device('cpu')))
model.eval()
print('Network loading complete.')
cap = cv.VideoCapture(0)
cap.set(cv.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv.CAP_PROP_FRAME_HEIGHT, 480)
print('Camera ready.')
while cap.isOpened():
    ret, image_origin = cap.read()
    if not ret:
        break
    imgL = image_origin[:, :640, :]
    imgR = image_origin[:, 640:, :]
    image = image_origin.copy()
    imgL_x = imgL.transpose(2, 0, 1)/256
    imgR_x = imgR.transpose(2, 0, 1)/256
    with torch.no_grad():
        imgL_x = torch.tensor(imgL_x, dtype=torch.float32)
        imgR_x = torch.tensor(imgR_x, dtype=torch.float32)
        img = torch.stack((imgL_x, imgR_x), dim=0)
        predict = model(img)
    if classify:
        imgL = image_generate_conf(imgL, predict[0])
        imgR = image_generate_conf(imgR, predict[1])
    else:
        imgLpred = predict[0][0], predict[1][0], predict[2][0]
        imgRpred = predict[0][1], predict[1][1], predict[2][1]
        imgL = image_generate_loc(imgL, imgLpred)
        imgR = image_generate_loc(imgR, imgRpred)
    image[:, :640, :] = imgL
    image[:, 640:, :] = imgR
    cv.imshow('test', image)
    c = cv.waitKey(10)
    if c == ord('q'):
        break
    elif c == ord('s'):
        cv.imwrite('err.png', image)
cap.release()
